#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
启动 Labelme 并加载训练插件
交互式 GUI 应用

使用方法:
    conda activate labelme
    python start_labelme_with_plugin.py
"""

import sys
import os
from pathlib import Path

def find_labelme_python():
    """查找 labelme 环境中的 Python"""
    # 方法1: 从 CONDA_PREFIX 获取（当前激活的环境）
    conda_prefix = os.environ.get('CONDA_PREFIX', '')
    if conda_prefix and 'labelme' in conda_prefix:
        python_path = os.path.join(conda_prefix, 'bin', 'python')
        if os.path.exists(python_path):
            return python_path
    
    # 方法2: 从 CONDA_PREFIX 的父目录获取（如果当前在其他 conda 环境）
    if conda_prefix and 'envs' in conda_prefix:
        conda_base = os.path.dirname(conda_prefix)
        labelme_python = os.path.join(conda_base, 'envs', 'labelme', 'bin', 'python')
        if os.path.exists(labelme_python):
            return labelme_python
    
    # 方法3: 从环境变量获取
    conda_base = os.environ.get('CONDA_BASE', '')
    
    # 方法4: 尝试从常见路径获取
    if not conda_base:
        home = os.path.expanduser('~')
        common_bases = [
            os.path.join(home, 'miniconda3'),
            os.path.join(home, 'anaconda3'),
            '/opt/miniconda3',
            '/opt/anaconda3',
        ]
        for base in common_bases:
            labelme_python = os.path.join(base, 'envs', 'labelme', 'bin', 'python')
            if os.path.exists(labelme_python):
                return labelme_python
    
    # 方法5: 使用 conda info 命令获取（如果 conda 可用）
    if not conda_base:
        try:
            import subprocess
            result = subprocess.run(
                ['conda', 'info', '--base'],
                capture_output=True,
                text=True,
                timeout=5
            )
            if result.returncode == 0:
                conda_base = result.stdout.strip()
        except:
            pass
    
    if conda_base:
        labelme_python = os.path.join(conda_base, 'envs', 'labelme', 'bin', 'python')
        if os.path.exists(labelme_python):
            return labelme_python
    
    return None

# 检查是否在正确的环境中
current_env = os.environ.get('CONDA_DEFAULT_ENV', '')
python_path = sys.executable

# 检查 Python 是否在 labelme 环境中
# 不仅要检查环境变量，还要检查 Python 路径和能否导入 labelme
is_labelme_env = False

# 方法1: 检查 Python 路径是否包含 labelme
if 'labelme' in python_path:
    is_labelme_env = True

# 方法2: 尝试导入 labelme 来验证（最可靠的方法）
if not is_labelme_env:
    try:
        import labelme
        is_labelme_env = True
    except ImportError:
        pass

# 如果环境变量是 labelme 但 Python 路径不对，需要切换
if current_env == 'labelme' and not is_labelme_env:
    # 环境变量显示是 labelme，但 Python 解释器不对，需要切换
    is_labelme_env = False

# 如果不在 labelme 环境，尝试自动切换
if not is_labelme_env:
    # 尝试自动切换到 labelme 环境
    labelme_python = find_labelme_python()
    
    # 如果没找到，尝试直接使用常见路径
    if not labelme_python:
        home = os.path.expanduser('~')
        common_paths = [
            os.path.join(home, 'miniconda3', 'envs', 'labelme', 'bin', 'python'),
            os.path.join(home, 'anaconda3', 'envs', 'labelme', 'bin', 'python'),
        ]
        for test_path in common_paths:
            if os.path.exists(test_path):
                labelme_python = test_path
                break
    
    if labelme_python and os.path.exists(labelme_python):
        print("=" * 60)
        print("自动切换到 labelme 环境")
        print(f"使用 Python: {labelme_python}")
        print("=" * 60)
        print()
        # 重新使用 labelme 环境的 Python 执行
        # 注意：os.execv 会替换当前进程，如果成功不会返回
        try:
            os.execv(labelme_python, [labelme_python] + sys.argv)
        except Exception as e:
            print(f"切换失败: {e}")
            print("请手动运行: conda activate labelme")
            sys.exit(1)
        # 如果执行到这里，说明 execv 失败了（不应该到达这里）
        print("切换失败，请手动运行: conda activate labelme")
        sys.exit(1)
    else:
        print("=" * 60)
        print("❌ 错误: 无法找到 labelme 环境")
        print("=" * 60)
        print(f"当前 Python: {python_path}")
        print(f"当前环境: {current_env}")
        print()
        print("请运行以下命令安装环境:")
        print("  bash paddleocr/install/setup_labelme_plugin_env.sh")
        print()
        print("然后使用以下方式启动:")
        print("  bash start_labelme_with_plugin.sh")
        print("  或")
        print("  conda activate labelme")
        print("  python start_labelme_with_plugin.py")
        print("=" * 60)
        sys.exit(1)

# 添加插件路径
plugin_path = Path(__file__).parent / 'labelme-plugin'
if str(plugin_path) not in sys.path:
    sys.path.insert(0, str(plugin_path))

try:
    # 导入 labelme
    from labelme import app
    
    # 导入插件
    from training_plugin import register_plugin
    
    # 修改 MainWindow 类以加载插件
    original_init = app.MainWindow.__init__
    
    def new_init(self, *args, **kwargs):
        original_init(self, *args, **kwargs)
        try:
            register_plugin(self)
            print("✓ 训练插件已加载")
        except Exception as e:
            print(f"⚠ 警告: 无法加载训练插件: {e}")
            import traceback
            traceback.print_exc()
    
    app.MainWindow.__init__ = new_init
    
    # 启动应用
    if __name__ == '__main__':
        print("=" * 60)
        print("启动 Labelme (带训练插件)")
        print("=" * 60)
        print("提示:")
        print("  - 在右侧面板找到 '训练管理' 面板")
        print("  - 点击 '🚀 开始训练' 按钮开始训练")
        print("  - 训练将在 seal_env 环境中执行（自动切换）")
        print("=" * 60)
        print()
        
        # 使用 QApplication 启动（labelme 的标准方式）
        from PyQt5.QtWidgets import QApplication
        qapp = QApplication(sys.argv)
        main_window = app.MainWindow()
        
        # 自动打开 paddleocr/dataset 目录
        try:
            # 获取项目根目录
            script_file = Path(__file__).resolve()
            project_root = script_file.parent.parent
            dataset_dir = project_root / 'paddleocr' / 'dataset'
            
            # 如果 dataset 目录存在，自动打开
            if dataset_dir.exists():
                # 优先打开 images 子目录（如果存在）
                images_dir = dataset_dir / 'images'
                if images_dir.exists() and images_dir.is_dir():
                    if hasattr(main_window, 'openDir'):
                        main_window.openDir(str(images_dir))
                        print(f"✓ 自动打开目录: {images_dir}")
                    elif hasattr(main_window, 'open_dir'):
                        main_window.open_dir(str(images_dir))
                        print(f"✓ 自动打开目录: {images_dir}")
                else:
                    # 如果没有 images 子目录，打开 dataset 目录
                    if hasattr(main_window, 'openDir'):
                        main_window.openDir(str(dataset_dir))
                        print(f"✓ 自动打开目录: {dataset_dir}")
                    elif hasattr(main_window, 'open_dir'):
                        main_window.open_dir(str(dataset_dir))
                        print(f"✓ 自动打开目录: {dataset_dir}")
            else:
                print(f"⚠ 数据集目录不存在: {dataset_dir}")
                print(f"   请手动打开目录或创建数据集目录")
        except Exception as e:
            print(f"⚠ 自动打开目录失败: {e}")
            import traceback
            traceback.print_exc()
        
        main_window.show()
        sys.exit(qapp.exec_())
        
except ImportError as e:
    print("=" * 60)
    print("❌ 错误: 无法导入 labelme 或插件")
    print("=" * 60)
    print(f"错误信息: {e}")
    print()
    print("请确保:")
    print("  1. 已安装 labelme: pip install labelme")
    print("  2. 已安装 PyQt5: pip install PyQt5")
    print("  3. 在 labelme 环境中运行: conda activate labelme")
    print()
    print("安装环境请运行:")
    print("  bash paddleocr/install/setup_labelme_plugin_env.sh")
    print("=" * 60)
    sys.exit(1)
except Exception as e:
    print(f"❌ 启动失败: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)


